package com.holly.unit.kafka.impl;

import com.holly.unit.kafka.config.KafkaConfig;
import com.holly.unit.kafka.model.KlocalThread;
import com.holly.unit.kafka.model.KproducerType;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.apache.kafka.clients.producer.KafkaProducer;
import org.apache.kafka.clients.producer.ProducerConfig;

import java.util.Properties;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicIntegerArray;
import java.util.concurrent.locks.ReentrantLock;

/**
 * 提供kafka生产者，线程池保证线程安全
 * @author yangpengpeng
 * 
 * 说明：
 * 为什么使用kafka连接池
 * 1：首先Apache Kafka 的所有通信都是使用TCP，在最初的创建连接的时候会，建立起TCP连接并且连接到整个Kafka集群，如下代码时，就会创建一个连接并且不会自动关闭。
 * 在此连接下调用KafkaProducer实例方法，即使共享我们并不改变这个KafkaProducer实例，只是通过这个KafkaProducer实例，发送消息，所以是线程安全的。
 * 2:不要创建太多TCP连接，TCP连接是要消耗系统资源的。由于发送数据是根据业务逻辑来进行的，可能当前TCP连接一直处于空闲状态了，则可以让其他线程使用发送，避免一直空闲
 * 3:为了保证分区的幂等性，一个线程发数据始终只绑定唯一的连接对象
 * 4:如果开启事务，则一个线程只对应一个数据源，如果不开事务，则可以多个线程对应一个数据源
 */
@Slf4j
@Data
public class KafkaProducerPool{
	private ReentrantLock  tranactionalLock = new ReentrantLock();
	public static ThreadLocal<KlocalThread> localVar = new ThreadLocal<KlocalThread>();
	public static final String CLIENT_ID = "holly";
	public static final String T_CLIENT_ID = "t_holly";

	private KafkaConfig kafkaConfig;
	private static final ConcurrentHashMap<String, KafkaProducer<String,Object>> connectionMap = new ConcurrentHashMap<String, KafkaProducer<String,Object>>();

	//事务的clientid
	private static int t_index = 0;
	//连接池初始化大小
	private int minPoolSize = 3;
	//连接池最大连接数
	private int maxPoolSize = 8;
	//活跃的连接数
	private int activeCount = 0;

	//连接对象数组
	private  KafkaProducer<String,Object>[] connections;
	//状态
	private AtomicIntegerArray states;

	@SuppressWarnings("unchecked")
	public KafkaProducerPool(KafkaConfig kafkaConfig) throws RuntimeException
	{	
		this.kafkaConfig = kafkaConfig;
		this.minPoolSize = kafkaConfig.producerMinPoolSize;
		this.maxPoolSize = kafkaConfig.producerMaxPoolSize;
		this.connections = new KafkaProducer[kafkaConfig.producerMaxPoolSize];
		this.states = new AtomicIntegerArray(new int[maxPoolSize]);

		log.info("连接池初始化,是否懒加载:{}",kafkaConfig.isLazy);

		if(kafkaConfig.isLazy) return;

		//初始化生产者最小的连接池
		for(int i = 0;i<minPoolSize;i++)
		{
			Properties props = createProperties(i);
			connections[i] = new KafkaProducer<String, Object>(props);
			activeCount++;
		}
	}

	private Properties createProperties(int i) {
		// 1. 创建用于连接Kafka的Properties配置
		Properties props = new Properties();
		props.put("bootstrap.servers",kafkaConfig.getBootstrapServers());
		props.put("acks", "all");
		props.put("key.serializer", kafkaConfig.keySerializer);
		props.put("value.serializer", kafkaConfig.valueSerializer);
		props.put("linger.ms",kafkaConfig.lingerMs);
		props.put("retries", kafkaConfig.Retries);
		props.put("batch.size", kafkaConfig.batchSize);
		props.put("batch.size", kafkaConfig.batchSize);
		props.put("enable.idempotence", true);
		props.put("client.id", CLIENT_ID+"_"+i);
		return props;
	}
	/**
	 * 从线程池中获取一个生产者
	 * @param timeOut 超时时间
	 * @return
	 */
	public KafkaProducer<String,Object> getProducer(long timeOut) throws RuntimeException
	{
		String tid =   localVar.get().getKeyId();
		if(tid !=null && connectionMap.get(tid) != null)
		{
			return connectionMap.get(tid);
		}
		//如果开启事务，如果事务id一致，可以用同一个producer,如果事务id不一致，则一个萝卜一个坑
		if(localVar.get().getType() == KproducerType.PRODUCER_TYPE_TRANCTION)
		{
			tranactionalLock.lock();
			try {
				String tranationalId = localVar.get().getKeyId();
				if(null == tranationalId) tranationalId = kafkaConfig.tranactionalId;
				Properties props = new Properties();
				props.put("bootstrap.servers",kafkaConfig.getBootstrapServers());
				props.put("acks", "all");
				props.put("key.serializer", kafkaConfig.keySerializer);
				props.put("value.serializer", kafkaConfig.valueSerializer);
				props.put(ProducerConfig.TRANSACTIONAL_ID_CONFIG,tranationalId);
				props.put("linger.ms",kafkaConfig.lingerMs);
				props.put("retries", kafkaConfig.Retries);
				props.put("batch.size", kafkaConfig.batchSize);
				props.put("client.id", T_CLIENT_ID+"_"+t_index++);
				props.put("enable.idempotence", true);
				KafkaProducer<String,Object> kafkaProducer = new KafkaProducer<String, Object>(props);
				connectionMap.put(tranationalId, kafkaProducer);
				kafkaProducer.initTransactions();
				return kafkaProducer;
			} catch (Exception e) {
				log.error("kafka分配事务源报错:{}",e);
			}finally
			{
				tranactionalLock.unlock();
			}
		}

		while(true)
		{
			//如果不开始事务，则没有状态。没有状态那就争抢吧。随便抢（只有第一次是随便抢，后面就绑定了， 1个producer 可以对应多个线程）
			//可能存在分配不均，这个后面考虑一下
			for(int i = 0;i<minPoolSize;i++)
			{
				if(states.get(i) == 0)
				{
					if(states.compareAndSet(i, 0, 1))
					{
						if(connections[i] == null)
						{
							Properties props = createProperties(i+activeCount);
							connections[i] = new KafkaProducer<String, Object>(props);
							activeCount++;
						}
						connectionMap.put(tid, connections[i]);
						return connections[i];
					}
				}
			}
			synchronized (this) {
				try {
					this.wait(10);
				} catch (Exception e) {
					throw new RuntimeException(e);
				}
			}
		}
	}
	/**
	 * 释放生产者，方便线程争抢
	 * @param producer
	 */
	public void freeProducer(KafkaProducer<String,Object> producer)
	{
		for(int i = 0;i<maxPoolSize;i++)
		{
			if(connections[i] == producer)
			{
				if(states.compareAndSet(i, 1, 0))
				{
					synchronized (this) {
						this.notifyAll();
					}
				}

			}
		}
	}
	/**
	 * 删除生产者
	 * @param producer
	 */
	public void removeProducer()
	{
		String tid =  localVar.get().getKeyId();
		KafkaProducer<String,Object> producer =  connectionMap.remove(tid);
		localVar.remove();
		producer.close();
		producer = null;
	}
	
}
